#!/usr/bin/env python
import dpdata
import matplotlib.pyplot as plt
import numpy as np

# 读取训练数据点
training_systems = dpdata.LabeledSystem("/home/zxg/be-cu_dp_course/dpgen/run/deepmd_file/Be16_hcp0.96_1.04/deepmd", fmt = "deepmd/npy")
predict = training_systems.predict("/home/zxg/be-cu_dp_course/dpgen/run/iter.000000/00.train/000/frozen_model.pb")

# 提取训练和预测的力分量
train_forces = training_systems["forces"]
predict_forces = predict["forces"]
train_energies = training_systems["energies"]
predict_energies = predict["energies"]

# 提取 x、y、z 轴力的分量
train_forces_x = [force[0] for sublist in train_forces for force in sublist]
train_forces_y = [force[1] for sublist in train_forces for force in sublist]
train_forces_z = [force[2] for sublist in train_forces for force in sublist]

predict_forces_x = [force[0] for sublist in predict_forces for force in sublist]
predict_forces_y = [force[1] for sublist in predict_forces for force in sublist]
predict_forces_z = [force[2] for sublist in predict_forces for force in sublist]

# 计算差值
energydiff = np.array(train_energies) - np.array(predict_energies)
forcediff_x = np.array(train_forces_x) - np.array(predict_forces_x)
forcediff_y = np.array(train_forces_y) - np.array(predict_forces_y)
forcediff_z = np.array(train_forces_z) - np.array(predict_forces_z)

# 创建一个figure
plt.figure(figsize=(12, 9))  # 设置figure的大小

# 对差值进行排序
sorted_indices_energy = np.argsort(energydiff)
sorted_indices_x = np.argsort(forcediff_x)
sorted_indices_y = np.argsort(forcediff_y)
sorted_indices_z = np.argsort(forcediff_z)

# 提取排序后的差值数组
sorted_energydiff = energydiff[sorted_indices_energy]
sorted_forcediff_x = forcediff_x[sorted_indices_x]
sorted_forcediff_y = forcediff_y[sorted_indices_y]
sorted_forcediff_z = forcediff_z[sorted_indices_z]

# 绘制能量差值图
plt.subplot(2, 2, 1)
plt.plot(sorted_energydiff, color='blue', label='Energy Difference (DFT - DeepMD)')
plt.title('Energy Difference')
plt.xlabel('Data Points')
plt.ylabel('Energy Difference (DFT - DeepMD)')
plt.legend()

# 绘制 x 轴力差值图
plt.subplot(2, 2, 2)
plt.plot(sorted_forcediff_x, color='green', label='Forces Difference x (DFT - DeepMD)')
plt.title('Forces Difference x')
plt.xlabel('Data Points')
plt.ylabel('Forces Difference x (DFT - DeepMD)')
plt.legend()

# 绘制 y 轴力差值图
plt.subplot(2, 2, 3)
plt.plot(sorted_forcediff_y, color='red', label='Forces Difference y (DFT - DeepMD)')
plt.title('Forces Difference y')
plt.xlabel('Data Points')
plt.ylabel('Forces Difference y (DFT - DeepMD)')
plt.legend()

# 绘制 z 轴力差值图
plt.subplot(2, 2, 4)
plt.plot(sorted_forcediff_z, color='blue', label='Forces Difference y (DFT - DeepMD)')
plt.title('Forces Difference y')
plt.xlabel('Data Points')
plt.ylabel('Forces Difference y (DFT - DeepMD)')
plt.legend()

# 设置图片的分辨率和保存路径
plt.tight_layout()  # 调整子图之间的间距
plt.savefig('inter00_f_e_direct_value.png', dpi=300)